"""
A2A协议类型定义 - 基于官方示例简化版本
"""
from datetime import datetime
from enum import Enum
from typing import Any, Literal, List, Optional, Union, Dict
from uuid import uuid4
from pydantic import BaseModel, Field, field_serializer


class TaskState(str, Enum):
    """任务状态枚举"""
    SUBMITTED = 'submitted'
    WORKING = 'working'
    INPUT_REQUIRED = 'input-required'
    COMPLETED = 'completed'
    CANCELED = 'canceled'
    FAILED = 'failed'
    UNKNOWN = 'unknown'


class TextPart(BaseModel):
    """文本部分"""
    kind: Literal['text'] = 'text'  # 使用A2A标准的kind字段
    text: str
    metadata: Optional[dict] = None


class FilePart(BaseModel):
    """文件部分"""
    kind: Literal['file'] = 'file'  # 使用A2A标准的kind字段
    file: dict
    metadata: Optional[dict] = None


class DataPart(BaseModel):
    """数据部分"""
    kind: Literal['data'] = 'data'  # 使用A2A标准的kind字段
    data: dict
    metadata: Optional[dict] = None


# 联合类型
Part = Union[TextPart, FilePart, DataPart]


class Message(BaseModel):
    """消息"""
    messageId: str = Field(default_factory=lambda: uuid4().hex)
    role: Literal['user', 'agent']
    parts: List[Part]
    contextId: Optional[str] = None
    taskId: Optional[str] = None
    kind: Literal['message'] = 'message'
    extensions: Optional[List[str]] = None
    referenceTaskIds: Optional[List[str]] = None
    metadata: Optional[dict] = None


class MessageSendConfiguration(BaseModel):
    """消息发送配置"""
    acceptedOutputModes: List[str] = ['text']
    blocking: Optional[bool] = None
    historyLength: Optional[int] = None
    pushNotificationConfig: Optional[dict] = None


class MessageSendParams(BaseModel):
    """A2A标准消息发送参数"""
    id: Optional[str] = Field(default_factory=lambda: uuid4().hex)
    message: Message
    configuration: Optional[MessageSendConfiguration] = None
    metadata: Optional[dict] = None


class TaskStatus(BaseModel):
    """任务状态"""
    state: TaskState
    message: Optional[Message] = None
    timestamp: datetime = Field(default_factory=datetime.now)

    @field_serializer('timestamp')
    def serialize_dt(self, dt: datetime, _info):
        return dt.isoformat()


class Artifact(BaseModel):
    """工件"""
    name: Optional[str] = None
    description: Optional[str] = None
    parts: List[Part]
    metadata: Optional[dict] = None
    index: int = 0
    append: Optional[bool] = None
    lastChunk: Optional[bool] = None


class Task(BaseModel):
    """任务 - 符合A2A协议标准"""
    id: str
    contextId: str = Field(..., description="上下文ID，用于跨交互的上下文对齐")
    status: TaskStatus
    history: Optional[List[Message]] = None
    artifacts: Optional[List[Artifact]] = None
    metadata: Optional[dict] = None
    kind: Literal['task'] = Field(default='task', description="固定值，标识为任务对象")
    # 保留sessionId以向后兼容
    sessionId: Optional[str] = None


class TaskStatusUpdateEvent(BaseModel):
    """任务状态更新事件 - 符合A2A协议标准"""
    taskId: str = Field(..., description="任务ID")
    contextId: str = Field(..., description="上下文ID") 
    kind: Literal['status-update'] = Field(default='status-update', description="事件类型")
    status: TaskStatus = Field(..., description="任务状态")
    final: bool = Field(default=False, description="是否为最终状态")
    metadata: Optional[Dict[str, Any]] = Field(default=None, description="事件元数据")


class TaskArtifactUpdateEvent(BaseModel):
    """
    任务工件更新事件
    Sent by server during sendStream or subscribe requests
    """

    append: bool | None = None
    """
    Indicates if this artifact appends to a previous one
    """
    artifact: Artifact
    """
    Generated artifact
    """
    contextId: str
    """
    The context the task is associated with
    """
    kind: Literal['artifact-update'] = 'artifact-update'
    """
    Event type
    """
    lastChunk: bool | None = None
    """
    Indicates if this is the last chunk of the artifact
    """
    metadata: dict[str, Any] | None = None
    """
    Extension metadata.
    """
    taskId: str
    """
    Task id
    """

# JSON-RPC消息类型
class JSONRPCMessage(BaseModel):
    """JSON-RPC消息基类"""
    jsonrpc: Literal['2.0'] = '2.0'
    id: Optional[Union[int, str]] = Field(default_factory=lambda: uuid4().hex)


class JSONRPCRequest(JSONRPCMessage):
    """JSON-RPC请求"""
    method: str
    params: Optional[dict] = None


class JSONRPCError(BaseModel):
    """JSON-RPC错误"""
    code: int
    message: str
    data: Optional[Any] = None


class JSONRPCResponse(JSONRPCMessage):
    """JSON-RPC响应"""
    result: Optional[Any] = None
    error: Optional[JSONRPCError] = None


# 任务相关参数
class TaskSendParams(BaseModel):
    """发送任务参数"""
    id: str
    sessionId: Optional[str] = Field(default_factory=lambda: uuid4().hex)
    message: Message
    acceptedOutputModes: Optional[List[str]] = None
    pushNotification: Optional[dict] = None
    historyLength: Optional[int] = None
    metadata: Optional[dict] = None


class TaskQueryParams(BaseModel):
    """任务查询参数"""
    id: str
    historyLength: Optional[int] = None
    metadata: Optional[dict] = None


# Agent相关类型
class AgentCapabilities(BaseModel):
    """智能体能力"""
    streaming: bool = False
    pushNotifications: bool = False
    stateTransitionHistory: bool = False


class AgentSkill(BaseModel):
    """智能体技能"""
    id: str
    name: str
    description: Optional[str] = None
    tags: Optional[List[str]] = None
    examples: Optional[List[str]] = None
    inputModes: Optional[List[str]] = None
    outputModes: Optional[List[str]] = None


class AgentProvider(BaseModel):
    """智能体提供者"""
    organization: str
    url: Optional[str] = None


class AgentCard(BaseModel):
    """智能体卡片"""
    name: str
    description: Optional[str] = None
    url: str
    provider: Optional[AgentProvider] = None
    iconUrl: Optional[str] = None
    version: str
    documentationUrl: Optional[str] = None
    capabilities: AgentCapabilities
    authentication: Optional[dict] = None
    defaultInputModes: List[str] = ['text']
    defaultOutputModes: List[str] = ['text']
    skills: List[AgentSkill]


# 请求响应类型
class SendTaskRequest(JSONRPCRequest):
    """发送任务请求"""
    method: Literal['tasks/send'] = 'tasks/send'
    params: TaskSendParams


class SendTaskResponse(JSONRPCResponse):
    """发送任务响应"""
    result: Optional[Task] = None


class GetTaskRequest(JSONRPCRequest):
    """获取任务请求"""
    method: Literal['tasks/get'] = 'tasks/get'
    params: TaskQueryParams


class GetTaskResponse(JSONRPCResponse):
    """获取任务响应"""
    result: Optional[Task] = None


# 错误类型
class TaskNotFoundError(JSONRPCError):
    """任务未找到错误"""
    code: int = -32001
    message: str = 'Task not found'
    data: None = None


class InvalidRequestError(JSONRPCError):
    """无效请求错误"""
    code: int = -32600
    message: str = 'Request payload validation error'
    data: Optional[Any] = None


class InternalError(JSONRPCError):
    """内部错误"""
    code: int = -32603
    message: str = 'Internal error'
    data: Optional[Any] = None


class MethodNotFoundError(JSONRPCError):
    """方法未找到错误"""
    code: int = -32601
    message: str = 'Method not found'
    data: None = None


# 定义 IframeSecurity 模型
class IframeSecurity(BaseModel):
    allowedOrigins: Optional[list[str]] = Field(None, description="允许的源")
    csp: Optional[str] = Field(None, description="内容安全策略")
    requireAuth: Optional[bool] = Field(None, description="是否需要认证")
    tokenPassthrough: Optional[bool] = Field(None, description="是否传递认证令牌")

# 定义 IframeData 模型
class IframeData(BaseModel):
    iframeId: str = Field(..., description="iframe唯一标识")
    title: Optional[str] = Field(None, description="iframe标题")
    src: str = Field(..., description="iframe源URL")
    width: Optional[Union[str, int]] = Field(None, description="宽度")
    height: Optional[Union[str, int]] = Field(None, description="高度")
    sandbox: Optional[str] = Field(None, description="沙箱属性")
    allowFullscreen: Optional[bool] = Field(None, description="是否允许全屏")
    allowPayment: Optional[bool] = Field(None, description="是否允许支付")
    referrerPolicy: Optional[str] = Field(None, description="引用策略")
    loading: Optional[str] = Field(None, description="加载策略")#'lazy', 'eager'
    interactionMode: Optional[str] = Field(None, description="交互模式")#'view-only', 'interactive', 'form-submit'
    security: Optional[IframeSecurity] = Field(None, description="安全配置")
    metadata: Optional[dict] = Field(None, description="iframe元数据")